{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bd892c00-c656-47fc-b94d-de3f6ffe8b2b",
   "metadata": {},
   "source": [
    "# Mistral Ai Chat - ChatMistralAI\n",
    "\n",
    "https://python.langchain.com/docs/integrations/chat/mistralai\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6aea0c7a-d27a-410e-8f86-a96a3807218c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: langchain-mistralai\n",
      "Version: 0.1.0\n",
      "Summary: An integration package connecting Mistral and LangChain\n",
      "Home-page: https://github.com/langchain-ai/langchain\n",
      "Author: \n",
      "Author-email: \n",
      "License: MIT\n",
      "Location: /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages\n",
      "Requires: httpx, httpx-sse, langchain-core, tokenizers\n",
      "Required-by: \n"
     ]
    }
   ],
   "source": [
    "!pip show langchain_mistralai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c5a6562-ae3f-41e7-a096-8ba6517449f5",
   "metadata": {},
   "source": [
    "## Create ChatMistralAI "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5527ef90-8632-4d5c-ae03-29b38a97f00d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import utility\n",
    "from langchain_mistralai.chat_models import ChatMistralAI\n",
    "\n",
    "chat_model = ChatMistralAI(mistral_api_key=utility.get_env(\"MISTRAL_API_KEY\",None), model=\"open-mistral-7b\")\n",
    "\n",
    "# class ChatMistralAI(BaseChatModel):\n",
    "#     \"\"\"A chat model that uses the MistralAI API.\"\"\"\n",
    "\n",
    "#     client: MistralClient = Field(default=None)  #: :meta private:\n",
    "#     async_client: MistralAsyncClient = Field(default=None)  #: :meta private:\n",
    "#     mistral_api_key: Optional[SecretStr] = None\n",
    "#     endpoint: str = DEFAULT_MISTRAL_ENDPOINT\n",
    "#     max_retries: int = 5\n",
    "#     timeout: int = 120\n",
    "#     max_concurrent_requests: int = 64\n",
    "\n",
    "#     model: str = \"mistral-small\"\n",
    "#     temperature: float = 0.7\n",
    "#     max_tokens: Optional[int] = None\n",
    "#     top_p: float = 1\n",
    "#     \"\"\"Decode using nucleus sampling: consider the smallest set of tokens whose\n",
    "#        probability sum is at least top_p. Must be in the closed interval [0.0, 1.0].\"\"\"\n",
    "#     random_seed: Optional[int] = None\n",
    "#     safe_mode: bool = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c427e7c-9334-4d33-926c-979cf41cc306",
   "metadata": {},
   "source": [
    "## Create Retriever backed by LanceDB vectorstore\n",
    "\n",
    "Data has already been embeded into LanceDB, please refer to `ollama-embeding.ipynb` for detail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6d751631-d27e-4c50-b05c-dea0b637d4d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 230 ms, sys: 44.4 ms, total: 275 ms\n",
      "Wall time: 671 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "import lancedb\n",
    "from langchain_community.vectorstores import LanceDB\n",
    "\n",
    "db = lancedb.connect(\"lancedb\")\n",
    "# connect to the existing table\n",
    "language = db.open_table('language')\n",
    "\n",
    "language_vectorstore = LanceDB(language, utility.get_embeddings())\n",
    "\n",
    "retriever = language_vectorstore.as_retriever(search_kwargs={\"k\": 1})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c8e292-7b1e-48db-ade9-ae7becec7750",
   "metadata": {},
   "source": [
    "## Create context retrieval for keyword"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4184632a-eee3-4ce4-8f86-c9071e3f0dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnableParallel, RunnablePassthrough\n",
    "from langchain_core.prompts import format_document\n",
    "\n",
    "DEFAULT_DOCUMENT_PROMPT = ChatPromptTemplate.from_template(template=\"{page_content}\")\n",
    "\n",
    "def _extract_context_documents(\n",
    "    docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
    "):\n",
    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
    "    context_content = document_separator.join(doc_strings)\n",
    "\n",
    "    if utility.is_trace_enabled() :\n",
    "        utility.trace(f'context length ({len(context_content)})')\n",
    "        utility.trace(f'<context>\\n{context_content}\\n</context>\\n')\n",
    "    \n",
    "    return context_content\n",
    "\n",
    "# retrieval = RunnableParallel(\n",
    "#     {\"context\": retriever, \"keyword\": RunnablePassthrough()}\n",
    "# )\n",
    "# retrieval.invoke(\"workspace\")\n",
    "\n",
    "# from operator import itemgetter\n",
    "# retrieval = RunnableParallel(\n",
    "#     {\"context\": itemgetter(\"keyword\") | retriever | _extract_context_documents, \"keyword\": itemgetter(\"keyword\") }\n",
    "# )\n",
    "# retrieval.invoke({\"keyword\":\"workspace\"})\n",
    "\n",
    "retrieval = RunnableParallel(\n",
    "    {\"context\": RunnablePassthrough() | retriever | _extract_context_documents, \"keyword\": RunnablePassthrough()}\n",
    ")\n",
    "# retrieval.invoke(\"workspace\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e4f47af-ed95-4c40-b685-83e64658cb1e",
   "metadata": {},
   "source": [
    "## Create OutputParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bbe0984a-cda8-4ca0-8137-1cfbe52edd9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "output_parser = StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0a0ff8b-7ae8-4331-ba5b-6f2077bb8181",
   "metadata": {},
   "source": [
    "## Create Prompt/Chat Message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ea0859b-300c-4da8-b15d-f24d581cb014",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 214 µs, sys: 19 µs, total: 233 µs\n",
      "Wall time: 273 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "template = \"\"\"\n",
    "You are an AI assistant with knowledge of structurizr DSL (Domain Specific Language) for software architecture modeling and documenation. \n",
    "\n",
    "Your task is to generate code snippet for keywords in structurizr DSL. \n",
    "\n",
    "<example>\n",
    "    <keyword>person</keyword>\n",
    "    <code>\n",
    "    ```\n",
    "    person <name> [description] [tags] \n",
    "    ```\n",
    "    </code>\n",
    "</example>\n",
    "\n",
    "<instruction>\n",
    "Extract code snippet for the keyword from the given context below precisely without any change. \n",
    "Keep the orignal code snippet format, do not generate explanations or examples.\n",
    "</instruction>\n",
    "\n",
    "<keyword>\n",
    "{keyword}\n",
    "</keyword>\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "736e983e-9aed-4168-8e3d-74bdbc6ea9d0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Create RAG Chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fce6e9f4-f42e-414b-a9d0-d6681686223b",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = retrieval | prompt | chat_model | output_parser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d9b7df9-ac35-4107-a246-365c0428ebdc",
   "metadata": {},
   "source": [
    "## Invoke chain to generate keyword code snippets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37dba09c-cb6f-4ab5-b441-9ad6198629ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```\n",
      "workspace [name] [description] {\n",
      "}\n",
      "\n",
      "workspace extends [file|url] {\n",
      "}\n",
      "```\n",
      "\n",
      "\n",
      "CPU times: user 42 ms, sys: 16.3 ms, total: 58.3 ms\n",
      "Wall time: 2.08 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"workspace\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b89f5a67-ff55-4f52-a335-f1f787377ddc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```\n",
      "container <name> [description] [technology] [tags] {\n",
      "  // content\n",
      "}\n",
      "```\n",
      "\n",
      "\n",
      "CPU times: user 24.6 ms, sys: 5.15 ms, total: 29.8 ms\n",
      "Wall time: 1.23 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"container\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "92b187c5-1e3d-4aeb-b765-1d7745c0898d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<code>\n",
      "component <name> [description] [technology] [tags] {\n",
      "...\n",
      "}\n",
      "</code>\n",
      "\n",
      "\n",
      "CPU times: user 24 ms, sys: 4.84 ms, total: 28.9 ms\n",
      "Wall time: 1.22 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"component\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d7b2fb6a-6e56-4d88-850c-7e3d722ef741",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```\n",
      "softwareSystem <name> [description] [tags] {\n",
      "...\n",
      "}\n",
      "```\n",
      "\n",
      "\n",
      "CPU times: user 23.8 ms, sys: 4.33 ms, total: 28.1 ms\n",
      "Wall time: 815 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"softwareSystem\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "28d53ca2-8694-4591-9c52-3e99b77c0573",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```\n",
      "deploymentNode <name> [description] [technology] [tags] [instances] {\n",
      "}\n",
      "```\n",
      "\n",
      "\n",
      "CPU times: user 22 ms, sys: 5 ms, total: 27 ms\n",
      "Wall time: 865 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"deploymentNode\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ffef58d2-e166-4add-8f0d-c780fec3d86c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<code>\n",
      "container <software system identifier> [description] {\n",
      "...\n",
      "}\n",
      "</code>\n",
      "\n",
      "\n",
      "CPU times: user 22.3 ms, sys: 3.93 ms, total: 26.3 ms\n",
      "Wall time: 1.14 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"container view\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a2e596f-d567-4851-84f4-c2be9cfc60e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"person\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "361ef7e7-4ea4-479a-9e67-d4432ea3adf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"model\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decb35b8-a381-42bb-868f-3ceb75b4df71",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "chat_response = chain.invoke(\"views\")\n",
    "\n",
    "utility.info(f'{chat_response}\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b931933-d79f-43e8-a8a4-511a16c71aa0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
